from common.clickhouse_alc_engine import (
    get_db,
    get_db_driver_client,
    get_factor_tb,
    get_kafka_db,
    get_kline_tb,
)
from common.tools import escape_param
from common.variables import TimeFrameType
import hashlib
import json


def get_watch_klines_last_topic(
    timeframe: TimeFrameType, symbols: list[str], name: str, group: str
) -> str:
    symbols = sorted(list(set(symbols)))
    hash_str = hashlib.md5(str(symbols).encode()).hexdigest()[:6]
    signal_str = f"{name}_{timeframe}_{len(symbols)}_{hash_str}"
    create_sql = """
     CREATE TABLE IF NOT EXISTS {kafka_tb} (
        time DateTime64(3),
        created_at DateTime DEFAULT now()
      ) ENGINE = Kafka SETTINGS kafka_broker_list = 'xx_kafka:9092',
                                kafka_topic_list = {topic},
                                kafka_group_name = {topic_group},
                                kafka_format = 'JSONEachRow',
                                kafka_num_consumers = 1;
    """

    view_sql = """
    CREATE MATERIALIZED VIEW IF NOT EXISTS {kafka_mv}
         TO {kafka_tb} (time)
    AS
    (
        select time
        from {kline_tb}
        where symbol in {symbols}
        group by time
        having count(distinct symbol) >= {len_symbols}
        order by time desc
        limit 1 );
    """
    kafka_tb = f"{get_kafka_db()}.{signal_str}"
    kafka_mv = f"{get_db()}.mv_{signal_str}"
    kline_tb = f"{get_db()}.{get_kline_tb(timeframe)}"
    client = get_db_driver_client(get_db())
    client.execute(
        create_sql.format(
            **{
                "kafka_tb": kafka_tb,
                "topic": escape_param(signal_str),
                "topic_group": escape_param(group),
            }
        )
    )
    client.execute(
        view_sql.format(
            **{
                "kafka_mv": kafka_mv,
                "kafka_tb": kafka_tb,
                "kline_tb": kline_tb,
                "symbols": escape_param(set(symbols)),
                "len_symbols": len(symbols),
            }
        )
    )
    return signal_str


def get_watch_factors_last_topic(
    timeframe: TimeFrameType,
    fnames: list[str],
    symbols: list[str],
    name: str,
    group: str,
) -> str:
    symbols = sorted(list(set(symbols)))
    fnames = sorted(list(set(fnames)))
    hash_str = hashlib.md5((str(symbols) + str(fnames)).encode()).hexdigest()[:6]
    signal_str = f"{name}_{timeframe}_{len(fnames)}_{len(symbols)}_{hash_str}"
    create_sql = """
         CREATE TABLE IF NOT EXISTS {kafka_tb} (
            time DateTime64(3),
            created_at DateTime DEFAULT now()
          ) ENGINE = Kafka SETTINGS kafka_broker_list = 'xx_kafka:9092',
                                    kafka_topic_list = {topic},
                                    kafka_group_name = {topic_group},
                                    kafka_format = 'JSONEachRow',
                                    kafka_num_consumers = 1;
        """

    view_sql = """
        CREATE MATERIALIZED VIEW IF NOT EXISTS {kafka_mv}
            TO {kafka_tb} (time)
        AS
        (
            select time
            from {factor_tb}
            where symbol in {symbols}
              and fname in {factors}
            group by time
            having count(distinct symbol, fname) >= {multi_cnt}
            order by time desc
            limit 1 );
        """
    kafka_tb = f"{get_kafka_db()}.{signal_str}"
    kafka_mv = f"{get_db()}.mv_{signal_str}"
    factor_tb = f"{get_db()}.{get_factor_tb(timeframe)}"
    client = get_db_driver_client(get_db())
    client.execute(
        create_sql.format(
            **{
                "kafka_tb": kafka_tb,
                "topic": escape_param(signal_str),
                "topic_group": escape_param(group),
            }
        )
    )
    client.execute(
        view_sql.format(
            **{
                "kafka_mv": kafka_mv,
                "kafka_tb": kafka_tb,
                "factor_tb": factor_tb,
                "symbols": escape_param(set(symbols)),
                "factors": escape_param(set(fnames)),
                "multi_cnt": len(symbols) * len(fnames),
            }
        )
    )
    return signal_str


def watch_kafka_topic(topic: str, group: str) -> dict:
    from kafka import KafkaConsumer

    # 创建 Kafka 消费者
    consumer = KafkaConsumer(
        topic,  # 订阅的 Kafka topic
        bootstrap_servers="xx_kafka:9092",  # Kafka 服务地址
        group_id=group,  # 消费者组 ID
        auto_offset_reset="latest",
    )

    try:
        for message in consumer:
            msg = message.value.decode("utf-8")
            yield json.loads(msg)
    finally:
        consumer.close()


async def async_watch_kafka_topic(topic: str, group: str) -> dict:
    # 创建 Kafka 消费者
    from aiokafka import AIOKafkaConsumer

    consumer = AIOKafkaConsumer(
        topic,  # 订阅的 Kafka topic
        bootstrap_servers="xx_kafka:9092",  # Kafka 服务地址
        group_id=group,  # 消费者组 ID
        auto_offset_reset="latest",
    )
    try:
        async for message in consumer:
            msg = message.value.decode("utf-8")
            yield json.loads(msg)
    finally:
        await consumer.stop()
